import torch
import torch.nn as nn

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()

        # 创建一个形状为 (max_len, d_model) 的位置编码矩阵
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # (max_len, 1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(torch.log(torch.tensor(10000.0)) / d_model))

        # 通过正弦和余弦函数计算位置编码
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数列
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数列

        pe = pe.unsqueeze(0)  # 增加 batch_size 维度 (1, max_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len]  # 获取前 seq_len 个位置编码，并与输入加和


class SingleTaskTransformer(nn.Module):
    def __init__(self, input_dim, task_dim, num_heads, num_layers, hidden_dim, dropout):
        super(SingleTaskTransformer, self).__init__()

        self.input_dim = input_dim
        self.task_dim = task_dim

        # 定义嵌入层
        self.embedding = nn.Linear(input_dim, hidden_dim)

        # 定义位置编码
        self.positional_encoding = PositionalEncoding(d_model=hidden_dim)

        # 创建多个 transformer encoder layer
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=dropout,
            batch_first=True  # 确保这里是 True
        )
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)

        # 定义单个输出层
        self.task_head = nn.Linear(hidden_dim, task_dim)

    def forward(self, x):
        # x shape: (batch_size, seq_length, input_dim)

        # 输入数据经过嵌入层
        x = self.embedding(x)  # shape: (batch_size, seq_length, hidden_dim)

        # 添加位置编码
        x = self.positional_encoding(x)  # shape: (batch_size, seq_length, hidden_dim)

        # 经过 transformer 编码器
        x = self.transformer_encoder(x)  # shape: (batch_size, seq_length, hidden_dim)

        # 对序列维度求均值
        x = x.mean(dim=1)  # shape: (batch_size, hidden_dim)

        # 经过单个输出层得到预测结果
        output = self.task_head(x)  # shape: (batch_size, task_dim)

        return output